# 滤掉某频率以上的信号
import numpy.fft as nf
import numpy as np
import matplotlib.pyplot as plt
import os
import linecache
import re


# 设定滤波（带阻）波段的上下限
lower_freq = 49   # 带阻波段上限

# 读取文件
folder = '电流电压信号/'  # 数据存储的文件夹名
files = os.listdir(folder)
for file in files:
    if file.startswith(".") or file.endswith("1") or file.endswith(".png") or file.endswith(".pdf"):  #
        # 过滤掉以“.”开头的文件和以“1”、“.png”、".pdf"结尾的文件
        continue

    data = np.loadtxt(f"{folder}{file}", skiprows=9)    # 跳过前9行数据
    t = data[:, 0]  # 第一列数据为时间序列
    f_t = data[:, 1]    # 第二列数据为函数序列

    fig, axes = plt.subplots(1, 2)  # 画两个子图，分别为时域图和频域图
    fig.dpi = 300   # 设置画出的图像的dpi，dpi越大，图像越清晰
    # 打印波形图
    axes[0].plot(t, f_t)
    axes[0].set_ylabel("signal")

    # 傅立叶变换
    F = nf.fft(f_t)
    N = f_t.size    # 采样点数
    delta_t = t[1] - t[0]     # 采样间隔
    f = nf.fftfreq(f_t.size, delta_t)   # 频域的横坐标序列
    mask = np.where(f > 0)  # 只分析频率大于零的部分（小于零的部分与之完全对称）
    A = 2*abs(F[mask]) / N  # 频谱图的纵坐标变换为实际振幅
    # 打印频谱图
    axes[1].plot(f[mask], A)
    axes[1].set_ylabel("$A$")
    fig.savefig(f"{folder}{file}.png")  # 将分析结果保存在以频率信息命名的png图片中

    if file.endswith("i"):
        # 如果文件以i结尾，进行滤波分析
        fig_n, axes_n = plt.subplots(1, 2)  # 再画两个子图，分别为时域图和频域图
        fig_n.dpi = 300  # 设置画出的图像的dpi，dpi越大，图像越清晰
        # 分析并输出实验结果
        # meas_freq = float(file.split('Hz')[0])     # 从文件名中读取激振器的频率值
        sign_freq = (f[f > 0])[np.argmax(abs(F[mask]))]       # 找到频谱图山振幅最大时对应的频率值
        amplitude = np.max(A)   # 最大振幅值
        print("分量最大的频率分量为：{:.2f} Hz, 对应的幅值为：{:.2f} V".format(sign_freq, amplitude))
        # devi_freq = sign_freq - meas_freq   # 绝对误差
        # real_devi = devi_freq / meas_freq   # 相对误差
        # print("偏差为：{:.2f} Hz（{:.2%}）".format(devi_freq, real_devi))
        # text = "f = {:.2f} Hz, A = {:.2f} V, error = {:.2f} Hz, {:.2%}".format(sign_freq, amplitude, devi_freq, real_devi)
        # plt.suptitle(text)  # 将误差作为图的标题，显示在图的上方

        # 滤波处理，并打印滤波之后的波形图
        '''
        根据不同的滤波模式，下式右边选择不同的表达式：
        如果是带阻，用 F * (abs(f) < lower_freq + abs(f) > upper_freq)；
        如果是带通，用 F * (abs(f) > lower_freq * abs(f) < upper_freq)；
        如果是低通，用 F * (abs(f) < lower_freq)；
        如果是高通，用 F * (abs(f) > upper_freq)。
        '''
        F_f = F * (abs(f) < lower_freq)    # 滤波之后的傅立叶变换值
        f_t_f = nf.ifft(F_f)  # 傅立叶逆变换
        axes_n[0].plot(t, f_t_f.real, color='red')  # 画出滤波之后的时域波形
        axes_n[0].set_xlabel("time (s)")
        axes_n[0].set_ylabel("signal")

        # 打印滤波之后的频谱图
        A_f = 2*abs(F_f[mask]) / N  # 频谱图的纵坐标变换为实际振幅
        axes_n[1].plot(f[mask], A_f, color='red')
        axes_n[1].set_xlabel("frequency (Hz)")
        axes_n[1].set_ylabel("$A$")
        text = f"{lower_freq} Hz Low Pass fft filter"
        plt.suptitle(text)  # 将误差作为图的标题，显示在图的上方
        # plt.tight_layout(rect=[0, 0, 1, 0.92])  # 设置子图布局，避免重叠
        fig_n.savefig(f"{folder}{file}_filtered{lower_freq} Hz.png")    # 将分析结果保存在以频率信息命名的png图片中
plt.show()
